syntax = "proto3";

package tensorflow;

import "tensorflow/tsl/profiler/protobuf/profiler_service.proto";

message NewProfileSessionRequest {
  ProfileRequest request = 1;
  // The place where we will dump profile data. We will normally use
  // MODEL_DIR/plugins/profile as the repository root.
  string repository_root = 2;
  repeated string hosts = 3;  // host or host:port, port will be ignored.
  string session_id = 4;
}

message NewProfileSessionResponse {
  // Auxiliary error_message.
  string error_message = 1;

  // Whether all hosts had returned a empty trace.
  bool empty_trace = 2;
}

message EnumProfileSessionsAndToolsRequest {
  string repository_root = 1;
}

message ProfileSessionInfo {
  string session_id = 1;
  // Which tool data is available for consumption.
  repeated string available_tools = 2;
}

message EnumProfileSessionsAndToolsResponse {
  // Auxiliary error_message.
  string error_message = 1;
  // If success, the returned sessions information are stored here.
  repeated ProfileSessionInfo sessions = 2;
}

message ProfileSessionDataRequest {
  // The place where we will read profile data. We will normally use
  // MODEL_DIR/plugins/profile as the repository root.
  string repository_root = 1;
  string session_id = 2;
  // Which host the data is associated. if empty, data from all hosts are
  // aggregated.
  string host_name = 5;
  // Which tool
  string tool_name = 3;
  // Tool's specific parameters. e.g. TraceViewer's viewport etc
  map<string, string> parameters = 4;
}

message ProfileSessionDataResponse {
  // Auxiliary error_message.
  string error_message = 1;

  // Output format. e.g. "json" or "proto" or "blob"
  string output_format = 2;

  // TODO(jiesun): figure out whether to put bytes or oneof tool specific proto.
  bytes output = 3;
}
////////////////////////////////////////////////////////////////////////////////
// ProfileAnalysis service provide entry point for profiling TPU and for
// serving profiled data to TensorBoard through GRPC
////////////////////////////////////////////////////////////////////////////////
service ProfileAnalysis {
  // Starts a profiling session, blocks until it completes.
  // TPUProfileAnalysis service delegate this to TPUProfiler service.
  // Populate the profiled data in repository, then return status to caller.
  rpc NewSession(NewProfileSessionRequest) returns (NewProfileSessionResponse) {
  }
  // Enumerate existing sessions and return available profile tools.
  rpc EnumSessions(EnumProfileSessionsAndToolsRequest)
      returns (EnumProfileSessionsAndToolsResponse) {}
  // Retrieve specific tool's data for specific session.
  rpc GetSessionToolData(ProfileSessionDataRequest)
      returns (ProfileSessionDataResponse) {}
}
